import os
import numpy as np
import tensorflow as tf
from prepocess.utils import get_file, get_batch

from nets.cnn_example import inference, losses, trainning, evaluation
from nets.mobilenet.mobilenet_v2 import mobilenet

IMG_W = 224  # resize图像，太大的话训练时间久
IMG_H = 224
BATCH_SIZE = 64  # 每个batch要放多少张图片
CAPACITY = 640  # 一个队列最大多少
MAX_STEP = 5000000  # 一般大于10K
learning_rate = 0.0001  # 一般小于0.0001
restore_step = 0
ckpt = {}

# 获取批次batch
train_dir = '../data_set/'  # 训练样本的读入路径
logs_train_dir = '../train_copy/'  # logs存储路径
train, train_label, val, val_label = get_file(train_dir)
# 训练数据及标签

train_batch, train_label_batch = get_batch(train, train_label, IMG_W, IMG_H, BATCH_SIZE, CAPACITY)
val_batch, val_label_batch = get_batch(val, val_label, IMG_W, IMG_H, BATCH_SIZE, CAPACITY)

# 训练操作定义

train_logits, _ = mobilenet(train_batch,
                            num_classes=6,
                            depth_multiplier=1.0)

train_loss = losses(train_logits, train_label_batch)
train_op = trainning(train_loss, learning_rate)
train_acc = evaluation(train_logits, train_label_batch, 'train_acc')
val_acc = evaluation(train_logits, val_label_batch, 'val_acc')

# 这个是log汇总记录
summary_op = tf.summary.merge_all()

# 产生一个会话
sess = tf.Session()
train_writer = tf.summary.FileWriter(logs_train_dir, sess.graph)
# 产生一个saver来存储训练好的模型

#
# exclude = ['Variable_12', 'Variable_13', 'Variable_11', 'Variable_5']
# variables_to_restore = slim.get_variables_to_restore(exclude=exclude)
#
# saver = tf.train.Saver(variables_to_restore)
saver = {}
if os.path.exists(logs_train_dir):
    ckpt = tf.train.get_checkpoint_state(logs_train_dir)
    print("loading {}".format(ckpt.model_checkpoint_path))
    # ckpt_prefix = saver.restore(sess, ckpt.model_checkpoint_path)
    restore_step = int(ckpt.model_checkpoint_path.split('/')[-1].split('-')[-1])


    saver = tf.train.import_meta_graph(ckpt.model_checkpoint_path+".meta")
    saver.restore(sess, tf.train.latest_checkpoint(logs_train_dir))

    sess.run(tf.global_variables_initializer())
    variable_names = [v.name for v in tf.trainable_variables()]
    values = sess.run(variable_names)
    for k, v in zip(variable_names, values):
        print("Variable: ", k)
        print("Shape: ", v.shape)
            # print(v)
else:
    saver = tf.train.Saver()
    sess.run(tf.global_variables_initializer())

# 队列监控
coord = tf.train.Coordinator()  # 设置多线程协调器
threads = tf.train.start_queue_runners(sess=sess, coord=coord)

# 进行batch的训练
try:
    # 执行MAX_STEP步的训练，一步一个batch
    for step in np.arange(restore_step,MAX_STEP):
        if coord.should_stop():
            break
        # 启动以下操作节点，有个疑问，为什么train_logits在这里没有开启？
        _, tra_loss, tra_acc, va_acc = sess.run([train_op, train_loss, train_acc, val_acc])

        # 每隔50步打印一次当前的loss以及acc，同时记录log，写入writer
        if step % 10 == 0:
            print('Step %d, train loss = %.2f, train accuracy = %.2f%% val accuracy = %.2f%% ' % (
                step, tra_loss, tra_acc * 100.0, va_acc * 100.0))
            summary_str = sess.run(summary_op)
            train_writer.add_summary(summary_str, step)
        checkpoint_path = os.path.join(logs_train_dir, 'mobilenet_v2_1.0_224')
        ckpt_prefix = saver.save(sess, checkpoint_path,global_step=step)
        print(" ")
        '''
        # 每隔100步，保存一次训练好的模型
        if (step + 1) == MAX_STEP:
            checkpoint_path = os.path.join(logs_train_dir, 'thing.ckpt')
            saver.save(sess, checkpoint_path, global_step=step)
        '''

except tf.errors.OutOfRangeError:
    print('Done training -- epoch limit reached')

finally:
    coord.request_stop()
coord.join(threads)
sess.close()

'''
y_conv = deep_CNN(train_batch)

# 定义损失函数
cross_entropy = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(labels=train_label_batch, logits=y_conv))
train_step = tf.train.AdamOptimizer(1e-4).minimize(cross_entropy)

# 训练验证
correct_prediction = tf.equal(tf.argmax(y_conv, 1), tf.argmax(train_label_batch, 1))
accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))  # 数据类型转换

print("start train")

sess = tf.Session()
saver = tf.train.Saver()
sess.run(tf.global_variables_initializer())
coord = tf.train.Coordinator()
threads = tf.train.start_queue_runners(sess=sess, coord=coord)
_, train_loss, train_accuracy = sess.run([train_step, correct_prediction, accuracy])
for i in range(1000):
    if i % 100 == 0:
        print('step:%d, training accuracy %g' % (i, train_accuracy))
    saver.save(sess, "input_data/some")
'''
